# utils/memory_system.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
记忆系统，负责存储、检索和管理智能体的记忆（包括短期和长期）
"""

import os
import json
import logging
import uuid
import re  # <--- Added missing import
from typing import Dict, List, Any, Optional, Tuple
from datetime import datetime, timedelta, timezone

try:
    from utils.llm_service import LLMService
except ImportError:
    LLMService = None  # type: ignore

logger = logging.getLogger("AI-MindCare-System-Memory")

DEFAULT_MEMORY_PATH = "data/mindcare_memory_web.json"  # Using separate file for web
SHORT_TERM_HOURS = 24


class MemorySystem:
    """记忆系统，管理用户信息、咨询记录、对话历史等"""

    def __init__(self,
                 persistence_path: str = None,
                 llm_service: Optional[LLMService] = None):
        self.persistence_path = persistence_path
        self.llm_service = llm_service
        self.memory_config = {
            "short_term_days": 30,  # 默认短期记忆保存天数
            "consolidation_interval_hours": 12  # 默认记忆整合间隔
        }
        self._memory = self._load_memory()
        self._memory.setdefault("users", {})
        self._memory.setdefault("metadata", {})

        logger.info(f"记忆系统初始化完成。持久化路径: {persistence_path or '无'}")
        if not llm_service:
            logger.warning("LLM 服务未提供给 MemorySystem，记忆归纳功能将不可用。")

    def _now_utc(self) -> datetime:
        return datetime.now(timezone.utc)

    def _load_memory(self) -> Dict[str, Any]:
        """从文件中加载数据"""
        if self.persistence_path and os.path.exists(self.persistence_path):
            try:
                with open(self.persistence_path, 'r', encoding='utf-8') as f:
                    logger.info(f"从 {self.persistence_path} 加载记忆...")
                    loaded_data = json.load(f)
                    loaded_data.setdefault("users", {})
                    loaded_data.setdefault("metadata", {})
                    return loaded_data
            except (json.JSONDecodeError, IOError) as e:
                logger.error(
                    f"加载记忆文件 '{self.persistence_path}' 失败: {e}，将使用空记忆。")
        return {
            "users": {},
            "metadata": {},
        }

    def _save_memory(self):
        """将数据保存到文件"""
        if self.persistence_path:
            try:
                os.makedirs(os.path.dirname(self.persistence_path),
                            exist_ok=True)
                with open(self.persistence_path, 'w', encoding='utf-8') as f:
                    json.dump(self._memory,
                              f,
                              indent=2,
                              ensure_ascii=False,
                              default=str)
                logger.info(f"记忆已保存到 {self.persistence_path}")
            except IOError as e:
                logger.error(f"保存记忆文件 '{self.persistence_path}' 失败: {e}")
            except TypeError as e:
                logger.error(f"序列化记忆时失败: {e}", exc_info=True)

    def _get_user_conversation_history(self, user_id: str):
        """获取用户的完整对话历史(字典类型)"""
        user_data = self._memory.get("users", {}).get(user_id, {})
        return user_data.get("conversations", {})

    def _update_long_term_memory(self, user_id: str, new_memory: dict):
        """更新用户的长期记忆(字典类型)"""
        users = self._memory.get("users", {})
        user_record = users.setdefault(user_id, {})
        existing_memory = user_record.setdefault("long_term_memory", {})
        if not isinstance(existing_memory, dict):
            existing_memory = {}
        existing_memory.update(new_memory)
        user_record["long_term_memory"] = existing_memory
        users[user_id] = user_record
        self._memory["users"] = users
        self._save_memory()
        logger.info(f"用户 {user_id} 的长期记忆已更新。")

    def _load_long_term_memory(self, user_id: str):
        """加载用户的长期记忆(字典类型)"""
        return self._memory.get("users",
                                {}).get(user_id,
                                        {}).get("long_term_memory", {})

    def _get_user_buffers(self):
        """初始化用户对话缓冲区"""
        users = self._memory.get("users", {})
        user_buffers = {}
        for user_id, user_data in users.items():
            user_buffers[user_id] = {}
            conversation_history = user_data.get("conversations", {})
            for conversation_id, conversation in conversation_history.items():
                messages = conversation.get("messages", [])
                user_buffers[user_id][conversation_id] = messages[
                    -20:]  # 保留最近的20条对话信息
        return user_buffers

    def _create_user_data(self, user_id: str) -> Dict[str, Any]:
        """创建用户数据"""
        users = self._memory.setdefault("users", {})
        if user_id not in users:
            logger.info(f"用户 {user_id} 不存在，创建新的记忆记录。")
            users[user_id] = {
                "basic_info": {},
                "long_term_memory": {},
                "conversations": {},
                "last_updated": self._now_utc().isoformat()
            }
            self._memory["users"] = users
            self._save_memory()
        return users[user_id]

    def _get_user_data(
            self,
            user_id: str,
            create_if_not_exists: bool = True) -> Optional[Dict[str, Any]]:
        """获取用户数据（支持自动创建）"""
        users = self._memory.setdefault("users", {})
        if user_id not in users and create_if_not_exists:
            self._create_user_data(user_id)
        user_data = users.get(user_id, {})
        # 确保必要字段存在
        user_data.setdefault("basic_info", {})
        user_data.setdefault("long_term_memory", {})
        user_data.setdefault("conversations", {})
        user_data.setdefault("last_updated", self._now_utc().isoformat())
        self._memory["users"] = users
        self._save_memory()
        return user_data

    def add_or_update_user_info(self, user_id: str, basic_info: Dict[str,
                                                                     Any]):
        """更新用户基本信息"""
        user_data = self._get_user_data(user_id, create_if_not_exists=True)
        if user_data:
            user_data["basic_info"].update(basic_info)
            user_data["last_updated"] = self._now_utc().isoformat()
            self._save_memory()

    def get_user_info(self, user_id: str) -> Optional[Dict[str, Any]]:
        """获取用户基本信息"""
        user_data = self._get_user_data(user_id, create_if_not_exists=False)
        return user_data.get("basic_info") if user_data else None

    def get_all_users(self) -> Dict[str, Dict[str, Any]]:
        """获取所有用户信息"""
        return {
            uid: data.get("basic_info", {})
            for uid, data in self._memory.get("users", {}).items()
        }

    def get_all_user_ids(self) -> List[str]:
        """获取所有用户ID"""
        return list(self._memory.get("users", {}).keys())

    def create_conversation_context(self,
                                    user_id: str,
                                    conversation_id: str,
                                    context: Dict[str, Any] = None):
        """创建会话上下文（包含默认初始化字段）"""
        if not user_id or not conversation_id:
            logger.error("创建上下文失败：未提供 user_id 或 conversation_id")
            return

        # 初始化上下文结构
        initial_context = {
            "last_updated": self._now_utc().isoformat(),
            "topic": "",
            "messages": [],
            "status": "active"
        }
        if context:
            initial_context.update(context)

        conversations = self._memory.get("users",
                                         {}).get(user_id,
                                                 {}).get("conversations", {})
        if conversation_id not in conversations:
            conversations[conversation_id] = initial_context
            self._memory["users"][user_id]["conversations"] = conversations
            self._save_memory()
            logger.info(f"成功创建会话上下文：{conversation_id}（用户：{user_id}）")
        else:
            logger.warning(f"会话上下文已存在：{conversation_id}（用户：{user_id}）")

    def get_conversation_context(self, user_id: str,
                                 conversation_id: str) -> Dict[str, Any]:
        """获取会话上下文（不存在时返回空字典）"""
        if not user_id or not conversation_id:
            logger.error("获取上下文失败：未提供 user_id 或 conversation_id")
            return {}
        return self._memory.get("users",
                                {}).get(user_id,
                                        {}).get("conversations",
                                                {}).get(conversation_id, {})

    def update_conversation_context(self, user_id: str, conversation_id: str,
                                    updates: Dict[str, Any]):
        """更新会话上下文"""
        if not user_id or not conversation_id:
            logger.error("更新上下文失败：未提供 user_id 或 conversation_id")
            return

        conversations = self._memory.get("users",
                                         {}).get(user_id,
                                                 {}).get("conversations", {})
        if conversation_id in conversations:
            conversations[conversation_id].update(updates)
            conversations[conversation_id]["last_updated_utc"] = self._now_utc(
            ).isoformat()
            self._save_memory()
            logger.debug(f"更新会话上下文：{conversation_id}（用户：{user_id}）")
        else:
            logger.error(f"会话上下文不存在：{conversation_id}（用户：{user_id}）")

    def delete_conversation_context(self, user_id: str, conversation_id: str):
        """删除会话上下文"""
        if not user_id or not conversation_id:
            logger.error("删除上下文失败：未提供 user_id 或 conversation_id")
            return

        conversations = self._memory.get("users",
                                         {}).get(user_id,
                                                 {}).get("conversations", {})
        if conversation_id in conversations:
            del conversations[conversation_id]
            self._save_memory()
            logger.info(f"成功删除会话上下文：{conversation_id}（用户：{user_id}）")
        else:
            logger.warning(f"会话上下文不存在：{conversation_id}（用户：{user_id}）")

    def generate_topic_for_conversation(self, user_id: str,
                                        conversation_id: str) -> str:
        """为会话生成摘要（需配置 LLM 服务）"""
        logger.info(f"开始为会话 {conversation_id}（用户：{user_id}）生成摘要...")
        try:
            # 获取对话消息
            conversation = self.get_conversation_context(
                user_id, conversation_id)
            messages = conversation.get("messages", [])
            if not messages:
                return "会话中没有消息，无法生成摘要。"

            # 构建提示词
            prompt = "\n".join([
                f"{msg['role']} ({msg['timestamp']}): {msg['content']}"
                for msg in messages
            ])
            summary_prompt = f"""
            请对以下对话生成主题摘要：
            
            对话内容：
            {prompt}
            
            摘要要求：
            你是一名记忆提取和摘要生成师，
           #目标：
           1. 调用此工作流时，基于{{input}}，生成以下两种围绕用户的记忆：
           ①用户画像信息，包括用户个人档案（姓名、性别、家乡等）和偏好信息；尽量使用“用户的职业是”、“用户喜欢/不喜欢”这样的陈述句。
           ②关键事件：以用户为中心，以“时间＋主体＋动作/副词等＋结果＋用户对事件的感受”的格式进行总结，例如“用户的小学升学考试十分顺利。”
           #输出结果
           一个包含了多条单个记忆的json列表。
           输出的生成记忆最终形式为json文本，包括
           {start_time}：本轮总结记忆源会话的开始时间 ；
           {end_time}：本轮总结记忆源会话的结束时间; 
           {lon_memo_content}：生成的单条记忆的中文内容;
           {lon_memo_type}：包括用户画像/关键事件两种类型；
           {row_num}:一个数字，为按照随机顺序赋予的标识数字。
        #注意事项：
        不同记忆分条存储，例如用户输入“我叫小张，我高考特别好”应分别存储为{记忆一： 用户叫小张}；{记忆二：用户的高考成绩特别好。}
        尽量根据上下文判断记忆的相似性和矛盾性，相同记忆进行合并，尽量减少记忆的重复。
        单条记忆不超过20个中文字符。
        同一轮生成的记忆的{start_time}{end_time}是一样的。

        
            """

            if not self.llm_service:
                logger.error("LLM 服务未初始化，无法生成摘要")
                return "服务暂不可用，请稍后再试。"

            # 调用大模型生成会话的主题
            topic = self.llm_service.generate_response(
                prompt=summary_prompt,
                system_message="你是专业的心理咨询摘要生成器，需用简洁语言提炼关键信息",
                max_tokens=500,
                temperature=0.5)
            logger.info(f"成功生成会话主题摘要：{conversation_id}（用户：{user_id}）")
            conversation["topic"] = topic.strip()
            self.update_conversation_context(user_id, conversation_id,
                                             {"topic": topic.strip()})
            self._save_memory()
            return topic

        except Exception as e:
            logger.error(f"生成摘要失败：{str(e)}", exc_info=True)
            return "生成摘要时发生错误，请检查日志"
